package com.olympus.hermes.gateway.config;

import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.server.resource.BearerTokenAuthenticationToken;
import org.springframework.security.oauth2.server.resource.BearerTokenError;
import org.springframework.security.web.server.authentication.ServerAuthenticationConverter;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static org.springframework.security.oauth2.server.resource.BearerTokenErrors.invalidRequest;
import static org.springframework.security.oauth2.server.resource.BearerTokenErrors.invalidToken;

/**
 * @author Mickey
 * @since 2020.08.27
 */
@Component
public class CustomTokenAuthenticationConverter implements ServerAuthenticationConverter {
    private static final Pattern AUTHORIZATION_PATTERN = Pattern.compile(
            "^Bearer (?<token>[a-zA-Z0-9-._~+/]+=*)$",
            Pattern.CASE_INSENSITIVE);

    private final boolean allowUriQueryParameter = false;

    @Override
    public Mono<Authentication> convert(ServerWebExchange exchange) {
        return Mono.fromCallable(() -> token(exchange.getRequest()))
                .map(token -> {
                    if (token.isEmpty()) {
                        BearerTokenError error = invalidTokenError();
                        throw new OAuth2AuthenticationException(error);
                    }
                    return new BearerTokenAuthenticationToken(new DefaultOAuth2AccessToken(token).getValue());
                });
    }

    private String token(ServerHttpRequest request) {
        String authorizationHeaderToken = resolveFromAuthorizationHeader(request.getHeaders());
        String parameterToken = request.getQueryParams().getFirst(OAuth2AccessToken.ACCESS_TOKEN.toLowerCase());
        if (authorizationHeaderToken != null) {
            if (parameterToken != null) {
                BearerTokenError error = invalidRequest("Found multiple bearer tokens in the request");
                throw new OAuth2AuthenticationException(error);
            }
            return authorizationHeaderToken;
        }
        else if (parameterToken != null && isParameterTokenSupportedForRequest(request)) {
            return parameterToken;
        }
        return null;
    }

    private static String resolveFromAuthorizationHeader(HttpHeaders headers) {
        String authorization = headers.getFirst(HttpHeaders.AUTHORIZATION);
        if (StringUtils.startsWithIgnoreCase(authorization, OAuth2AccessToken.BEARER_TYPE.toLowerCase())) {
            Matcher matcher = AUTHORIZATION_PATTERN.matcher(authorization);

            if (!matcher.matches() ) {
                BearerTokenError error = invalidTokenError();
                throw new OAuth2AuthenticationException(error);
            }

            return matcher.group("token");
        }
        return null;
    }

    private static BearerTokenError invalidTokenError() {
        return invalidToken("Bearer token is malformed");
    }

    private boolean isParameterTokenSupportedForRequest(ServerHttpRequest request) {
        return this.allowUriQueryParameter && HttpMethod.GET.equals(request.getMethod());
    }
}
